from typing import List

from redis import Redis

from .nodes.mcp import MCPNodes
from ..graph.base import BaseGraph
from ..graph.nodes import *
from ..llms import LLMConfig, InsCodeModel
from ..mcps import MCPClientManager
from ..models.graph import GraphRequest
from ..parser.xml_parser import XmlToolResultParser
from ..config.agents import LLMType
from loguru import logger


class MultimediaGraph(BaseGraph):
    """
    专业的图像制作AI助手，专注于图像生成、编辑和处理等核心功能。
    应用场景：图像创作、图像分析、图像编辑、视觉内容制作等所有与图像处理相关的工作。

    核心功能：
    - 文生图：根据文字描述生成高质量的图像、插图、艺术作品
    - 图生图：基于现有图片进行风格转换、内容修改、质量提升
    - 图像分析：深度分析图片内容、风格、构图、色彩等视觉元素
    - 图像编辑：对图像进行修改、优化、格式转换等处理
    - 创意实现：协助用户将创意想法转化为具体的图像作品
    - 多格式支持：支持多种图像尺寸和格式的生成与处理
    - 风格转换：在不同艺术风格间进行图像风格转换
    - 图像优化：提升图像质量、调整视觉效果

    专业特长：
    - 图像生成：AI绘画、概念图、数字艺术创作
    - 图像处理：风格迁移、图像增强、格式转换
    - 内容分析：图像识别、视觉元素解析、构图分析
    - 创意设计：插画创作、艺术作品生成

    注意：本助手专注于图像的生成、编辑和处理，不具备UI原型图、交互原型、线框图等界面设计功能。

    系统会根据用户输入的问题自动分析并设置以下参数：
    - title (必选): 基于用户问题生成的标题，能概括图像创作需求或处理意图，控制在20个字以内。
        - 对于图像生成请求，标题应体现创作主题，如"风景画生成"、"人物肖像创作"
        - 对于图像处理请求，标题应体现处理目标，如"图像风格转换"、"图片质量优化"
        - 标题应该清晰反映用户的图像制作或处理需求
        - 标题生成是必需的，不可省略
    """
    llm_vl: InsCodeModel = None
    llm_image: InsCodeModel = None

    def __init__(self,
                 llm_vl: InsCodeModel | LLMConfig = None,
                 llm_image: InsCodeModel | LLMConfig = None,
                 redis: Redis = None,
                 nodes: List[BaseNodes] = [],
                 graph_name="agent_multimedia_",
                 agent_template: str = "multimedia",
                 mcp_client: MCPClientManager = None,
                 **kwargs):
        super().__init__(graph_name=graph_name,
                         agent_template=agent_template,
                         redis=redis,
                         mcp_client=mcp_client,
                         **kwargs)
        nodes = nodes if nodes else []

        self.llm_vl = InsCodeModel.build(llm_vl if llm_vl else LLMConfig.build_from_type(LLMType.LLM_VL))
        self.llm_image = InsCodeModel.build(llm_image if llm_image else LLMConfig.build_from_type(LLMType.LLM_IMAGE))

        self.nodes = nodes if nodes else [
            AskUserNodes(),
            ImageNodes(llm_vl=self.llm_vl, llm_image=self.llm_image),
            FileNodes(),
            CompletionNodes()
        ]

        if mcp_client and mcp_client.is_available():
            self.nodes.append(MCPNodes(mcp_client))

        self.parser = XmlToolResultParser(tools=self.get_nodes_schema())

    async def get_system_template(self, request: GraphRequest, state=None, template_type: str = None):
        """
        Override base implementation to add multimedia-specific default values.
        """
        if "custom_instructions" not in request.kwargs:
            request.kwargs["custom_instructions"] = ""

        return await super().get_system_template(request=request, state=state, template_type=template_type)

    async def destroy(self):
        """
        Clean up multimedia-specific resources and call parent cleanup.
        """
        try:
            if self.llm_vl:
                await self.llm_vl.close()
                self.llm_vl = None

            if self.llm_image:
                await self.llm_image.close()
                self.llm_image = None

            await super().destroy()
        except Exception as e:
            logger.error(f"Error during multimedia graph cleanup: {e}")
